package com.ai.controller;


import com.ai.common.annotation.Agent;
import com.alibaba.dashscope.common.Message;
import com.alibaba.fastjson.JSON;
import jakarta.annotation.Resource;
import jakarta.servlet.http.HttpServletRequest;
import lombok.Data;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatModel;

import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.context.ApplicationContext;
import org.springframework.web.bind.annotation.*;
import reactor.core.publisher.Flux;

import java.io.ByteArrayOutputStream;
import java.io.InputStream;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;

@RestController
@RequestMapping("/ai")
class AIController {

    @Resource
    private ChatModel chatModel;



    @GetMapping("/generate")
    public Map generate(@RequestParam(value = "message", defaultValue = "小明的成绩是多少") String message) {
        ChatClient chatClient = ChatClient.create(chatModel);
        // 获取带有Agent注解的bean
//        Map<String, Object> beansWithAnnotation = applicationContext.getBeansWithAnnotation(Agent.class);
//        String[] functionBeanNames = new String[beansWithAnnotation.size()];
//        functionBeanNames = beansWithAnnotation.keySet().toArray(functionBeanNames);
        return Map.of("generation", chatClient.prompt().user(message).call().content());
    }

    @Resource
    private  ApplicationContext applicationContext;

    @GetMapping(value = "/generateStream", produces = "text/html;charset=UTF-8")
    public Flux<String> generateStream(@RequestParam(value = "message", defaultValue = "小明的成绩是多少") String message) {
        ChatClient chatClient = ChatClient.builder(chatModel).build();
        // 获取带有Agent注解的bean
        Map<String, Object> beansWithAnnotation = applicationContext.getBeansWithAnnotation(Agent.class);
        String[] functionBeanNames = new String[beansWithAnnotation.size()];
        functionBeanNames = beansWithAnnotation.keySet().toArray(functionBeanNames);
        Flux<String> content = chatClient.prompt().user(message).functions(functionBeanNames).stream().content();
        return content;
    }

    @GetMapping(value = "/generateStream2", produces = "text/html;charset=UTF-8")
    public Flux<String> generateStream2(@RequestParam(value = "message", defaultValue = "小明的成绩是多少") String message) {
        ChatClient chatClient = ChatClient.builder(chatModel).build();
        // 获取带有Agent注解的bean
        Map<String, Object> beansWithAnnotation = applicationContext.getBeansWithAnnotation(Agent.class);
        String[] functionBeanNames = new String[beansWithAnnotation.size()];
        functionBeanNames = beansWithAnnotation.keySet().toArray(functionBeanNames);
        Flux<String> content = chatClient.prompt().user(message).stream().content();
        return content;
    }

    @Data
    public static class ChatRequest {
            private Boolean stream;
            private String messages = "显示器不亮了怎么办";
    }

    @RequestMapping(value = "/chat", produces = "text/html;charset=UTF-8")
    public Flux<String> chat(ChatRequest chatRequest, HttpServletRequest request) {
        String requestBody = getRequestBody(request);
        String messages = JSON.parseObject(requestBody).getString("messages");
        List<Message> messageList = JSON.parseArray(messages, Message.class);
        List<org.springframework.ai.chat.messages.Message> messagess = new ArrayList<>();
        messageList.forEach(message -> {
            if ("user".equals(message.getRole())) {
                messagess.add(new UserMessage(message.getContent()));
            } else if ("assistant".equals(message.getRole())) {
                messagess.add(new AssistantMessage(message.getContent()));
            } else if ("system".equals(message.getRole())) {
                messagess.add(new SystemMessage(message.getContent()));
            }
        });
        Map<String, String[]> parameterMap = request.getParameterMap();
        ChatClient chatClient = ChatClient.builder(chatModel).build();
        // 获取带有Agent注解的bean
        Map<String, Object> beansWithAnnotation = applicationContext.getBeansWithAnnotation(Agent.class);
        String[] functionBeanNames = new String[beansWithAnnotation.size()];
        functionBeanNames = beansWithAnnotation.keySet().toArray(functionBeanNames);
        Flux<String> content = chatClient.prompt(new Prompt(messagess)).functions(functionBeanNames).stream().content();
        return content;
    }

    private String getRequestBody(HttpServletRequest request) {
        try {
            InputStream v_inputstream = request.getInputStream();
            ByteArrayOutputStream baos = new ByteArrayOutputStream();
            int x = 0;
            while ((x = v_inputstream.read()) != -1) {
                baos.write(x);
            }
            baos.flush();
            return new String(baos.toByteArray(), StandardCharsets.UTF_8);
        } catch (Exception e) {
            System.out.println(e);
        }
        return "";
    }
}
